package org.example;

import com.sun.jna.Pointer;
import org.example.library.WintunLibrary;
import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.IpSelector;

import java.io.IOException;
import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;


public class ConnectServer {
    public String host;
    public int port;
    public String ip;
    public Pointer pointer;

    private static final ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor();

    AsynchronousSocketChannel clientChannel;

    public ConnectServer(String host, int port) {
        this.host = host;
        this.port = port;
        connect();
    }

    /**
     * 连接服务器
     */
    private void connect() {
        try {
            clientChannel = AsynchronousSocketChannel.open();
            clientChannel.connect(new InetSocketAddress(host, port)).get();
            var buff = ByteBuffer.allocate(100);
            var len = clientChannel.read(buff).get();
            if (len > 0) {
                buff.flip();
                int l = buff.getInt();
                var bt = new byte[l];
                buff.get(bt);
                //获取IP
                ip = new String(bt);
                //读任务
                executorService.submit(this::readTask);

            } else {
                clientChannel.close();
            }
        } catch (Exception e) {
            e.printStackTrace();
            if (clientChannel != null) {
                try {
                    clientChannel.close();
                } catch (IOException ex) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 上传数据到服务器
     * @param data
     */
    public void uploadData(byte[] data) {

        if (clientChannel != null && clientChannel.isOpen()) {
            try {
                var buff = ByteBuffer.allocate(data.length + 4);
                buff.putInt(data.length);
                buff.put(data);
                buff.flip();
                clientChannel.write(buff).get();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 读取服务器下发的数据
     */
    private void readTask() {
        try {
            var buff = ByteBuffer.allocate(4);
            while (!Thread.interrupted()) {
                var len = clientChannel.read(buff).get();
                if (len <= 0) {
                    break;
                }
                buff.flip();
                var anInt = buff.getInt();
                var buff_dy = ByteBuffer.allocate(anInt);
                var len_dy = clientChannel.read(buff_dy).get();
                if (len_dy <= 0) {
                    break;
                }
                var bt = new byte[len_dy];
                buff_dy.flip();
                buff_dy.get(bt);

                //写入网络设配器
                try {

                    IpPacket packet = (IpPacket) IpSelector.newPacket(bt, 0, bt.length);
                    writeData(packet.getRawData());

                } catch (Exception e) {
                    e.printStackTrace();
                }

                buff.clear();
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (clientChannel != null) {
                try {
                    clientChannel.close();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
        }
    }

    /**
     * 数据写到网络适配器
     * @param data
     */
    private void writeData(byte[] data) {

        try {
            IpPacket packet = (IpPacket) IpSelector.newPacket(data, 0, data.length);
            InetAddress dstAddr = packet.getHeader().getDstAddr();
            InetAddress srcAddr = packet.getHeader().getSrcAddr();


            if (dstAddr instanceof Inet4Address) {

                ByteBuffer allocate = ByteBuffer.allocate(data.length + 4);
                allocate.putInt(data.length);
                allocate.put(data);
                allocate.flip();

                Pointer packetPointer = WintunLibrary.INSTANCE.WintunAllocateSendPacket(pointer, data.length);

                if (packetPointer != null) {

                    packetPointer.write(0, data, 0, data.length);

                    // 发送数据包
                    WintunLibrary.INSTANCE.WintunSendPacket(pointer, packetPointer);

                } else {
                    System.err.println("Failed to allocate send packet.");
                }

            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
